#!/usr/bin/python


class QuickUnionPathCompression(object):

  def __init__(self, num_objects):
    self.num_objects = num_objects
    self.id_list = range(self.num_objects)

  def GetRoot(self, object_id):
    while not self.id_list[object_id] == object_id:
      # Point each object on the path to root to its grandparent.
      current_parent = self.id_list[object_id]
      self.id_list[object_id] = self.id_list[current_parent]
      object_id = self.id_list[object_id]
    return object_id

  def connected(self, p, q):
    return self.GetRoot(p) == self.GetRoot(q)

  def union(self, p, q):
    print 'Connecting %d and %d' % (p, q)
    p_root = self.GetRoot(p)
    q_root = self.GetRoot(q)
    self.id_list[p_root] = q_root
    print 'pos: ', range(self.num_objects)
    print 'ids: ', self.id_list

def main():
  num_objects = 10
  qup = QuickUnionPathCompression(num_objects)
  qup.union(1, 3)
  qup.union(2, 4)
  qup.union(1, 4)
  qup.union(1, 9)
  print '1, 9 connected: ', qup.connected(1, 9)
  print '1, 4 connected: ', qup.connected(1, 4)

if __name__ == '__main__':
  main() 
